#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
PDF拆分工具
功能：将PDF文件拆分为两个文件
- 001.pdf: 包含第一页
- 002.pdf: 包含剩余页面
"""

import os
import sys
from pathlib import Path
from PyPDF2 import PdfReader, PdfWriter


def create_output_directory(output_path):
    """
    创建输出目录，如果不存在的话
    
    Args:
        output_path (str): 输出目录路径
    """
    try:
        Path(output_path).mkdir(parents=True, exist_ok=True)
        print(f"✓ 输出目录已准备: {output_path}")
    except Exception as e:
        print(f"✗ 创建输出目录失败: {e}")
        sys.exit(1)


def split_pdf(input_pdf_path, output_dir):
    """
    拆分PDF文件
    
    Args:
        input_pdf_path (str): 输入PDF文件路径
        output_dir (str): 输出目录路径
    """
    try:
        # 检查输入文件是否存在
        if not os.path.exists(input_pdf_path):
            print(f"✗ 输入文件不存在: {input_pdf_path}")
            return False
        
        # 读取PDF文件
        print(f"📖 正在读取PDF文件: {input_pdf_path}")
        reader = PdfReader(input_pdf_path)
        total_pages = len(reader.pages)
        print(f"📄 PDF总页数: {total_pages}")
        
        if total_pages == 0:
            print("✗ PDF文件为空")
            return False
        
        # 创建输出目录
        create_output_directory(output_dir)
        
        # 拆分第一页
        if total_pages >= 1:
            print("📝 正在创建001.pdf (第一页)...")
            writer_first = PdfWriter()
            writer_first.add_page(reader.pages[0])
            
            output_file_001 = os.path.join(output_dir, "001.pdf")
            with open(output_file_001, 'wb') as output_file:
                writer_first.write(output_file)
            print(f"✓ 001.pdf 创建成功: {output_file_001}")
        
        # 拆分剩余页面
        if total_pages > 1:
            print("📝 正在创建002.pdf (剩余页面)...")
            writer_remaining = PdfWriter()
            for page_num in range(1, total_pages):
                writer_remaining.add_page(reader.pages[page_num])
            
            output_file_002 = os.path.join(output_dir, "002.pdf")
            with open(output_file_002, 'wb') as output_file:
                writer_remaining.write(output_file)
            print(f"✓ 002.pdf 创建成功: {output_file_002}")
        else:
            print("ℹ️  只有一页，不创建002.pdf")
        
        print("🎉 PDF拆分完成！")
        return True
        
    except Exception as e:
        print(f"✗ PDF拆分失败: {e}")
        return False


def main():
    """主函数"""
    # 配置参数
    input_pdf = "/Users/xueguang/Desktop/45f9b4823ac19a94fe89e3b76caa371f.pdf"
    output_directory = "./output"
    
    print("🚀 PDF拆分工具启动")
    print(f"📁 输入文件: {input_pdf}")
    print(f"📁 输出目录: {output_directory}")
    print("-" * 50)
    
    # 执行拆分
    success = split_pdf(input_pdf, output_directory)
    
    if success:
        print("\n✅ 任务完成！")
        print(f"📂 输出文件位置: {os.path.abspath(output_directory)}")
    else:
        print("\n❌ 任务失败！")
        sys.exit(1)


if __name__ == "__main__":
    main()
